use core::mem;

use alloc::vec::Vec;
use strum::{AsRefStr, Display, EnumString};

#[derive(Clone, Copy, Debug, PartialEq, Eq, AsRefStr, EnumString, Display)]
pub enum DType {
    F16,
    F32,
    F64,
    I8,
    I16,
    I32,
    I64,
    U8,
    U16,
    U32,
    U64,
    Bool,
    BF16,
}

impl DType {
    pub fn size_of(&self) -> usize {
        match self {
            DType::F32 | DType::I32 | DType::U32 => 4,
            DType::F64 | DType::I64 | DType::U64 => 8,
            DType::F16 | DType::I16 | DType::U16 => 2,
            DType::I8 | DType::U8 | DType::Bool => 1,
            DType::BF16 => 2,
        }
    }
}

#[derive(Clone)]
pub struct Tensor {
    pub dtype: DType,
    pub shape: Vec<usize>,
    pub data: Vec<u8>,
}

impl Tensor {
    pub fn new(dtype: DType, shape: Vec<usize>) -> Self {
        let count: usize = shape.iter().product();
        let bytes = count * dtype.size_of();
        Self {
            dtype,
            shape,
            data: alloc::vec![0; bytes],
        }
    }

    pub fn from_data<T: Copy>(dtype: DType, shape: &[usize], values: &[T]) -> Self {
        let count: usize = shape.iter().product();
        assert_eq!(
            count,
            values.len(),
            "Tensor::from_data: shape {:?} vs {} values",
            shape,
            values.len()
        );

        let mut data = Vec::with_capacity(count * mem::size_of::<T>());
        let p = unsafe {
            core::slice::from_raw_parts(values.as_ptr() as *const u8, count * mem::size_of::<T>())
        };
        data.extend_from_slice(p);

        Self {
            dtype,
            shape: shape.to_vec(),
            data,
        }
    }

    pub fn from_f32(shape: &[usize], values: Vec<f32>) -> Self {
        Self::from_data(DType::F32, shape, &values)
    }

    pub fn elem_count(&self) -> usize {
        self.shape.iter().product()
    }

    pub fn as_slice<T: Copy>(&self) -> &[T] {
        // # Safety: caller must ensure dtype matches T
        unsafe { core::slice::from_raw_parts(self.data.as_ptr() as *const T, self.elem_count()) }
    }

    pub fn as_mut_slice<T: Copy>(&mut self) -> &mut [T] {
        unsafe {
            core::slice::from_raw_parts_mut(self.data.as_mut_ptr() as *mut T, self.elem_count())
        }
    }
}
